package ordered_table

func NewSBTree() *SBTree {
	return &SBTree{Root: nil}
}

type CompareAble interface{
	CompareTo (CompareAble) int
	GetValue ()             int
}

type SBTreeNode struct {
	Key   CompareAble
	Value interface{}
	Size  int
	Left  *SBTreeNode
	Right *SBTreeNode
}

type SBTree struct {
	Root *SBTreeNode
}


func (s *SBTree) RightRotate(cur *SBTreeNode) *SBTreeNode {
	leftNode := cur.Left
	cur.Left = leftNode.Right
	leftNode.Right = cur
	leftNode.Size = cur.Size
	s.ReCalcCurSize(cur)
	return leftNode
}

func (s *SBTree) LeftRotate(cur *SBTreeNode) *SBTreeNode {
	rightNode := cur.Right
	cur.Right = rightNode.Left
	rightNode.Left = cur
	rightNode.Size = cur.Size
	s.ReCalcCurSize(cur)
	return rightNode
}

func (s *SBTree)ReCalcCurSize(cur *SBTreeNode) {
	cur.Size = 1
	if cur.Left != nil {
		cur.Size += cur.Left.Size
	}
	if cur.Right != nil {
		cur.Size += cur.Right.Size
	}
}

func (s *SBTree)maintain(cur *SBTreeNode) *SBTreeNode {
	if cur == nil {
		return nil
	}
	leftSize := s.GetNodeLeftSize(cur)
	leftLeftSize := s.GetNodeLeftLeftSize(cur)
	leftRightSize := s.GetNodeLeftRightSize(cur)
	rightSize := s.GetNodeRightSize(cur)
	rightRightSize := s.GetNodeRightRightSize(cur)
	rightLeftSize := s.GetNodeRightLeftSize(cur)
	if rightSize < leftLeftSize {
		cur = s.RightRotate(cur)
		cur.Right = s.maintain(cur.Right)
		cur = s.maintain(cur)
	} else if rightSize < leftRightSize {
		cur.Left = s.LeftRotate(cur.Left)
		cur = s.RightRotate(cur)
		cur.Left = s.maintain(cur.Left)
		cur.Right = s.maintain(cur.Right)
		cur = s.maintain(cur)
	} else if leftSize < rightRightSize {
		cur = s.LeftRotate(cur)
		cur.Left = s.maintain(cur.Left)
		cur = s.maintain(cur)
	} else if leftSize < rightLeftSize {
		cur.Right = s.RightRotate(cur.Right)
		cur = s.LeftRotate(cur)
		cur.Left = s.maintain(cur.Left)
		cur.Right = s.maintain(cur.Right)
		cur = s.maintain(cur)
	}
	return cur
}

func(s *SBTree) FindLastIndex(key CompareAble) *SBTreeNode {
	pre := s.Root
	cur := s.Root
	for cur != nil {
		pre = cur
		if key.CompareTo(cur.Key) == 0 {
			break
		} else if key.CompareTo(cur.Key) < 0 {
			cur = cur.Left
		} else {
			cur = cur.Right
		}
	}
	return pre
}

func(s *SBTree) FindLastNoSmallIndex(key CompareAble) *SBTreeNode {
	var ans *SBTreeNode
	cur := s.Root
	for cur != nil {
		if key.CompareTo(cur.Key) == 0 {
			break
		} else if key.CompareTo(cur.Key) < 0 {
			ans = cur
			cur = cur.Left
		} else {
			cur = cur.Right
		}
	}
	return ans
}

func(s *SBTree) FindLastNoBigIndex(key CompareAble) *SBTreeNode {
	var ans *SBTreeNode
	cur := s.Root
	for cur != nil {
		if key.CompareTo(cur.Key) == 0 {
			break
		} else if key.CompareTo(cur.Key) < 0 {
			cur = cur.Left
		} else {
			ans = cur
			cur = cur.Right
		}
	}
	return ans
}

func (s *SBTree) Add(cur *SBTreeNode, key CompareAble, value interface{}) *SBTreeNode {
	if cur == nil {
		return &SBTreeNode{Key: key, Value: value, Size: 1}
	} else {
		cur.Size++
		if key.CompareTo(cur.Key) < 0 {
			cur.Left = s.Add(cur.Left, key, value)
		} else {
			cur.Right = s.Add(cur.Right, key, value)
		}
		return s.maintain(cur)
	}
}

func (s *SBTree) Delete(cur *SBTreeNode, key CompareAble) *SBTreeNode {
	cur.Size--
	if key.CompareTo(cur.Key) > 0 {
		cur.Right = s.Delete(cur.Right, key)
	} else if key.CompareTo(cur.Key) < 0 {
		cur.Left = s.Delete(cur.Left, key)
	} else {
		if cur.Left == nil && cur.Right == nil {
			cur = nil
		} else if cur.Left == nil && cur.Right != nil {
			cur = cur.Right
		} else if cur.Left != nil && cur.Right == nil {
			cur = cur.Left
		} else {
			var pre *SBTreeNode
			des := cur.Right
			des.Size--
			for des.Left != nil {
				pre = des
				des = des.Left
				des.Size--
			}
			if pre != nil {
				pre.Left = des.Right
				des.Right = cur.Right
			}
			des.Left = cur.Left
			des.Size = des.Left.Size + s.GetNodeRightSize(des)
			cur = des
		}
	}
	return cur
}

func (s *SBTree)GetIndex(cur *SBTreeNode, kth int) *SBTreeNode {
	if cur == nil {return nil}
	if kth == s.GetNodeLeftSize(cur)+1 {
		return cur
	} else if kth <= s.GetNodeLeftSize(cur) {
		return s.GetIndex(cur.Left, kth)
	} else {
		return s.GetIndex(cur.Right, kth - s.GetNodeLeftSize(cur)-1)
	}
}

func (s *SBTree)Size() int {
	if s.Root == nil {return 0}
	return s.Root.Size
}

func (s *SBTree)ContainsKey(key CompareAble) bool {
	if key == nil {
		panic(any("key error"))
	}
	lastNode := s.FindLastIndex(key)
	if lastNode != nil && key.CompareTo(lastNode.Key) == 0 {
		return true
	}
	return false
}

func (s *SBTree)Put(key CompareAble, value interface{}) {
	if key == nil {
		panic(any("key error"))
	}
	lastNode := s.FindLastIndex(key)
	if lastNode != nil && key.CompareTo(lastNode.Key) == 0 {
		lastNode.Value = value
	} else {
		s.Root = s.Add(s.Root, key, value)
	}
}

func (s *SBTree)Remove(key CompareAble) {
	if key == nil {
		panic(any("key error"))
	}
	lastNode := s.FindLastIndex(key)
	if lastNode != nil && key.CompareTo(lastNode.Key) == 0 {
		s.Root = s.Delete(s.Root, key)
	}
}

func (s *SBTree)GetIndexKey(index int) CompareAble {
	if index < 0 || index >= s.Root.Size {
		panic(any("key error"))
	}
	return s.GetIndex(s.Root, index+1).Key
}

func (s *SBTree)GetIndexValue(index int) interface{} {
	if index < 0 || index >= s.Root.Size {
		panic(any("key error"))
	}
	return s.GetIndex(s.Root, index+1).Value
}

func (s *SBTree) Get (key CompareAble) interface{} {
	if key == nil {
		panic(any("key error"))
	}
	lastNode := s.FindLastIndex(key)
	if lastNode != nil && key.CompareTo(lastNode.Key) == 0 {
		return lastNode.Value
	}
	return nil
}

func (s *SBTree) FirstKey() CompareAble {
	if s.Root == nil {return nil}
	cur := s.Root
	for cur.Left != nil {
		cur = cur.Left
	}
	return cur.Key
}

func (s *SBTree) LastKey() CompareAble {
	if s.Root == nil {return nil}
	cur := s.Root
	for cur.Right != nil {
		cur = cur.Right
	}
	return cur.Key
}

func (s *SBTree) FloorKey(key CompareAble) CompareAble {
	if key == nil {
		panic(any("key error"))
	}
	lastNoBigNode := s.FindLastNoBigIndex(key)
	if lastNoBigNode == nil {
		return nil
	}
	return lastNoBigNode.Key
}

func (s *SBTree) CeilingKey(key CompareAble) CompareAble {
	if key == nil {
		panic(any("key error"))
	}
	lastNoSmallNode := s.FindLastNoSmallIndex(key)
	if lastNoSmallNode == nil {
		return nil
	}
	return lastNoSmallNode.Key
}

func(s *SBTree) GetNodeRightSize(node *SBTreeNode) int {
	if node.Right == nil {
		return 0
	}
	return node.Right.Size
}

func(s *SBTree) GetNodeRightRightSize(node *SBTreeNode) int {
	if node.Right == nil || node.Right.Right == nil {
		return 0
	}
	return node.Right.Right.Size
}

func(s *SBTree) GetNodeRightLeftSize(node *SBTreeNode) int {
	if node.Right == nil || node.Right.Left == nil {
		return 0
	}
	return node.Right.Left.Size
}

func(s *SBTree) GetNodeLeftSize(node *SBTreeNode) int {
	if node.Left == nil {
		return 0
	}
	return node.Left.Size
}

func(s *SBTree) GetNodeLeftLeftSize(node *SBTreeNode) int {
	if node.Left == nil || node.Left.Left == nil {
		return 0
	}
	return node.Left.Left.Size
}

func(s *SBTree) GetNodeLeftRightSize(node *SBTreeNode) int {
	if node.Left == nil || node.Left.Right == nil {
		return 0
	}
	return node.Left.Right.Size
}